#include "../../include/GlobalMemory.hpp"

__global__ void scan(int* data, int size, int* aux_array)
{
	extern __shared__ int v[];

	int tid = threadIdx.x;

	///load data into shared memory
	if(2*blockDim.x*blockIdx.x + tid < size)
		v[tid] = data[2*blockDim.x*blockIdx.x + tid];
	if(2*blockDim.x*blockIdx.x + tid + blockDim.x < size)
		v[tid + blockDim.x] = data[2*blockDim.x*blockIdx.x + tid + blockDim.x];
	__syncthreads();

	int step = 2;

	///phase 1: reduce
	for( ; step <= blockDim.x * 2; step *= 2)
	{
		if( tid*step + step - 1 < blockDim.x * 2)
			v[tid*step + step - 1] = v[tid*step + step - 1] + v[tid*step + step - 1 - step/2];

		__syncthreads();
	}

	///phase 1.5: store sum and zeroing last element
	if(tid == 0)
	{
		aux_array[blockIdx.x] = v[blockDim.x*2 - 1];

		v[blockDim.x*2 - 1] = 0;
	}
	
	///phase 2: sweep
	for(step /= 2; step >= 2; step /= 2)
	{
		if( tid*step + step - 1 < blockDim.x * 2)
		{
			int tmp = v[tid*step + step - 1]; ///get right

			///right = right + left
			v[tid*step + step - 1] = v[tid*step + step - 1] + v[tid*step + step - 1 - step/2];

			v[tid*step + step - 1 - step/2] = tmp; ///left = right
		}

		__syncthreads();
	}

	///set results back to global memory
	if(2*blockDim.x*blockIdx.x + tid < size)
		data[2*blockDim.x*blockIdx.x + tid] = v[tid];
	if(2*blockDim.x*blockIdx.x + tid + blockDim.x < size)
		data[2*blockDim.x*blockIdx.x + tid + blockDim.x] = v[tid + blockDim.x];
}

__global__ void final_add(int* data, int size, int* aux_array)
{
	__shared__ int v;

	if(threadIdx.x == 0)
		v = aux_array[blockIdx.x + 1];

	__syncthreads();

	int idx = (blockIdx.x + 1) * 2 * blockDim.x + threadIdx.x;
	if(idx < size)
		data[idx] = data[idx] + v;
	if(idx + blockDim.x < size)
		data[idx + blockDim.x] = data[idx + blockDim.x] + v;
}

void scan_on_gpu(CH::GlobalMemory<int>& i_d, const int& block_dim)
{
	///compute block #
	int block_num = (i_d.size() + 2*block_dim - 1) / (2*block_dim);

	///allocate aux array
	CH::GlobalMemory<int> aux_array( block_num );

	///load input from global into shared memory[implicit padding and continuous loading]
	///scan, fill aux array and write result back to global memory
	scan<<<block_num, block_dim, sizeof(int)*block_dim*2>>>(i_d.get_memory_ptr(), i_d.size(), aux_array.get_memory_ptr());

	///if block_num > 1
	if(block_num > 1)
	{
		scan_on_gpu( aux_array, block_dim );

		final_add<<<block_num - 1, block_dim>>>(i_d.get_memory_ptr(), i_d.size(), aux_array.get_memory_ptr());
	}
}
